import os
import glob
import json
import shutil
from safetensors.torch import load_file, save_file

def fix_safetensors_and_index_to_newdir(model_dir, tar_dir=None, prefix="model-"):
    # 1. 目标路径
    if tar_dir is None:
        base_dir = os.path.dirname(os.path.abspath(model_dir.rstrip("/")))
        src_name = os.path.basename(os.path.abspath(model_dir.rstrip("/")))
        tar_dir = os.path.join(base_dir, f"fix_{src_name}")
    os.makedirs(tar_dir, exist_ok=True)
    print(f"Will save all fixed files to {tar_dir}")

    # 2. 复制非safetensors、非index.json文件
    for file in glob.glob(os.path.join(model_dir, "*")):
        if file.endswith(".safetensors") or file.endswith("index.json"):
            continue
        if os.path.isfile(file):
            shutil.copy(file, tar_dir)
        elif os.path.isdir(file):
            # 如果有子目录，如 tokenizer, config，递归拷贝
            shutil.copytree(file, os.path.join(tar_dir, os.path.basename(file)), dirs_exist_ok=True)

    # 3. 修正所有 safetensors 文件
    safetensor_files = sorted(glob.glob(os.path.join(model_dir, f"{prefix}*.safetensors")))
    print(f"Found {len(safetensor_files)} safetensors files.")
    for path in safetensor_files:
        print(f"Processing {path} ...")
        tensors = load_file(path)
        new_tensors = {}
        for k, v in tensors.items():
            if k.startswith("model."):
                new_tensors[k[6:]] = v
            else:
                new_tensors[k] = v
        out_path = os.path.join(tar_dir, os.path.basename(path))
        save_file(new_tensors, out_path)
        print(f"Saved to {out_path}")

    # 4. 修正 index.json
    index_path = os.path.join(model_dir, "index.json")
    if os.path.exists(index_path):
        with open(index_path, 'r', encoding='utf-8') as f:
            index_data = json.load(f)
        new_weight_map = {}
        for k, v in index_data["weight_map"].items():
            new_k = k[6:] if k.startswith("model.") else k
            new_weight_map[new_k] = v
        index_data["weight_map"] = new_weight_map
        out_index_path = os.path.join(tar_dir, "index.json")
        with open(out_index_path, 'w', encoding='utf-8') as f:
            json.dump(index_data, f, ensure_ascii=False, indent=2)
        print(f"New index.json written to {out_index_path}")
    else:
        print("index.json not found, skipped index update.")

    print(f"\n所有权重、index及其他文件已拷贝/修正到 {tar_dir}。")
    print("用新的目录进行推理即可！")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="安全修正 Qwen2.5-VL merge 权重的 model. 前缀，输出到新目录")
    parser.add_argument("--dir", type=str, default='./experiments/med_pub', help="原safetensors和index.json目录")
    parser.add_argument("--tar_dir", type=str, default=None, help="输出目录（默认为 fix_前缀的新目录）")
    args = parser.parse_args()
    fix_safetensors_and_index_to_newdir(args.dir, args.tar_dir)
